import time
import numpy
import open3d as o3d
import os
from utils.MQ.message import send_message
import json


def judge_folder(path):
    if not os.path.exists(path):
        os.makedirs(path)


class NpyConvert():
    def __init__(self, MQ_send, connection_send, task_id, npy_file, target_folder, out_type='pcd', RADIO=80):
        self.task_id = task_id
        self.MQ_send = MQ_send
        self.connection_send = connection_send
        self.npy_file = npy_file
        self.target_folder = target_folder
        self.out_type = out_type
        self.RADIO = RADIO

    def npy2p_x_f(self):
        """
        npy to pcd or ply
        :param npy_file:
        :param out_type:
        :param out_dir:
        :return:
        """
        print("--------------- start output pcd")
        start = time.time()
        # long running
        # do something other
        z_points = numpy.load(self.npy_file)
        y_num, x_num = z_points.shape

        unit_x_point = numpy.arange(x_num)
        x_points = numpy.tile(unit_x_point, y_num) / self.RADIO

        unit_y_point = numpy.arange(y_num)
        y_points = numpy.repeat(unit_y_point, x_num) / self.RADIO

        z_points = z_points.reshape((-1))

        nan_index = numpy.isnan(z_points)

        # remove nan points
        x_points = x_points[~nan_index]
        y_points = y_points[~nan_index]
        z_points = z_points[~nan_index]

        pcd = o3d.geometry.PointCloud()
        m3d = numpy.array([x_points, y_points, z_points]).T
        pcd.points = o3d.utility.Vector3dVector(m3d)
        # file_name = osp.basename(npy_file).split('.')[0] + '.' + out_type
        pcd_file = os.path.basename(self.npy_file)[:-3] + self.out_type
        file_name = os.path.join(self.target_folder, pcd_file)
        judge_folder(os.path.dirname(file_name))
        o3d.io.write_point_cloud(file_name, pcd, write_ascii=True, compressed=False)
        end = time.time()
        print("pcd path: " + file_name)
        print("saving pcd done, used time: " + str(end - start))
        res = {
            "task_id": self.task_id,
            "task_type": "NPY CONVERT TO PCD",
            "pcd_path": file_name
        }
        send_message(self.MQ_send, self.connection_send, json.dumps(res, ensure_ascii=False))
        # return file_name


if __name__ == "__main__":
    ROOT = "D:\\Programs\\data\\dataset\\pcd_generate"
    FILES = os.listdir(ROOT)
    for FILE in FILES:
        if FILE[-3:]=='npy':
            con = NpyConvert(0, 0, 0, os.path.join(ROOT, FILE), ROOT, )
            try:
                con.npy2p_x_f()
            except Exception as e:
                print(e)
